from typing import Dict, Any, Optional

import jieba.posseg
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.callbacks import CallbackManagerForChainRun
from langchain_core.runnables import RunnableConfig

from ai_engine.car_wrap.prompts import EXTRACTION_PROMPT
from ai_engine.common.ai_common import trace_context
from ai_engine.core.base.base_chain import BaseChain
from ai_engine.core.constant.Area import get_area_by_name
from ai_engine.core.model.chat import ExtractionRequest


class ExtractionChain(BaseChain):
    request: ExtractionRequest

    def _call(
            self,
            inputs: Dict[str, Any],
            run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        question = inputs["question"]
        # 设置调用跟踪request_id
        if self.request.request_id:
            trace_context.set(self.request.request_id)
        # docs = self.knowledge_vs.similarity_search(question, self.request, _run_manager)
        output: Dict[str, Any] = {}
        seg_list = jieba.posseg.lcut(question.upper())
        # 同义词
        synonyms = []
        province_city_text = ""
        for s in seg_list:
            if s.flag in ['city', 'province']:
                name, code = get_area_by_name(s.word)
                if name:
                    if province_city_text:
                        province_city_text += "\n" + name + "|" + code
                    else:
                        province_city_text = name + "|" + code
        # 设置上下文参数
        new_inputs = inputs.copy()
        new_inputs["province_city"] = province_city_text
        new_inputs["context"] = []
        qa_chain = create_stuff_documents_chain(self.llm, EXTRACTION_PROMPT.PROMPT)
        answer = qa_chain.invoke(input=new_inputs,
                                 config=RunnableConfig(callbacks=_run_manager.get_child()))
        output[self.output_key] = answer
        return output
